// Date:   Thu Aug 14 07:27:00 PM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#include <openvino/runtime/tensor.hpp>
#include <vector>

#include "models/yolo.hpp"
#include "common.hpp"

static float IoU(const YoloBox& a, const YoloBox& b) {
  const float a_x2 = a.x + a.w;
  const float a_y2 = a.y + a.h;
  const float b_x2 = b.x + b.w;
  const float b_y2 = b.y + b.h;

  const float inter_x1 = std::max(a.x, b.x);
  const float inter_y1 = std::max(a.y, b.y);
  const float inter_x2 = std::min(a_x2, b_x2);
  const float inter_y2 = std::min(a_y2, b_y2);

  const float inter_area = std::max(0.f, inter_x2 - inter_x1) * std::max(0.f, inter_y2 - inter_y1);
  const float area_a = a.w * a.h;
  const float area_b = b.w * b.h;

  return inter_area / (area_a + area_b - inter_area);
}

static std::vector<YoloBox> NMS(const std::vector<YoloBox>& boxes, float iou_threshold) {
  std::vector<YoloBox> result;
  std::vector<bool> removed(boxes.size(), false);

  for (size_t i = 0; i < boxes.size(); ++i) {
    if (removed[i]) continue;
    result.push_back(boxes[i]);
    for (size_t j = i + 1; j < boxes.size(); ++j) {
      if (removed[j]) continue;
      if (boxes[i].class_id == boxes[j].class_id && IoU(boxes[i], boxes[j]) > iou_threshold) {
        removed[j] = true;
      }
    }
  }
  return result;
}

void yolov5_preprocess(cv::Mat const& img, int width, int height, ov::Tensor& out_tensor) {
  float* tensor_data = out_tensor.data<float>();

  const float r = std::min((float)width / img.cols, (float)height / img.rows);
  const int new_w = static_cast<int>(std::round(img.cols * r));
  const int new_h = static_cast<int>(std::round(img.rows * r));

  // 计算padding
  const int top = (height - new_h) / 2;
  const int left = (width - new_w) / 2;

  // resize后直接存放到临时Mat
  cv::Mat resized;
  cv::resize(img, resized, cv::Size(new_w, new_h));

  // 直接遍历batch_tensor对应内存，逐像素赋值
  // 先填充背景色(114/255)
  const float fill_val = 114.0f / 255.0f;
  for (int c = 0; c < 3; ++c) {
    float* ptr = tensor_data + c * height * width;
    for (int j = 0; j < height * width; ++j) {
      ptr[j] = fill_val;
    }
  }

  // 把 resized 图像数据写入 batch_tensor（BGR→RGB 顺序转换，归一化）
  for (int y = 0; y < new_h; ++y) {
    const uchar* src_row = resized.ptr<uchar>(y);
    for (int x = 0; x < new_w; ++x) {
      int dst_y = y + top;
      int dst_x = x + left;
      // batch_tensor 里是 CHW 排布，索引公式：
      // index = channel * H * W + y * W + x
      float* r_ptr = tensor_data + 0 * height * width + dst_y * width + dst_x;
      float* g_ptr = tensor_data + 1 * height * width + dst_y * width + dst_x;
      float* b_ptr = tensor_data + 2 * height * width + dst_y * width + dst_x;

      // BGR → RGB，归一化
      *r_ptr = src_row[3 * x + 2] / 255.0f;
      *g_ptr = src_row[3 * x + 1] / 255.0f;
      *b_ptr = src_row[3 * x + 0] / 255.0f;
    }
  }
}

std::vector<YoloBox> yolov5_postprocess(
  ov::Tensor& out_tensor,
  int width,
  int height,
  int input_width,
  int input_height,
  float conf_threshold,
  float nms_threshold
)
{
  auto const shape = out_tensor.get_shape();
  auto const batch_size = shape[0];
  assert(batch_size == 1);
  // num_det 表示所有特征图检测框的总数
  auto const num_det = shape[1];
  assert(shape.size() == 3);
  assert(batch_size == shape[0]);

  auto const M = shape[2];
  assert(M > 5);

  float* data = out_tensor.data<float>();

  std::vector<YoloBox> boxes;

  auto const scale_x = static_cast<float>(width) / input_width;
  auto const scale_y = static_cast<float>(height) / input_height;

  std::vector<YoloBox> pre_boxes;
  for (size_t i = 0; i < num_det; i++) {
    const float box_confidence = data[i * M + 4];
    if (box_confidence < conf_threshold)
        continue;

    // find class id with max confidence
    float max_class_prob = 0.f;
    int class_id = -1;
    for (size_t c = 5; c < M; c++) {
        float class_prob = data[i * M + c];
        if (class_prob > max_class_prob) {
            max_class_prob = class_prob;
            class_id = c - 5;
        }
    }

    const float confidence = box_confidence * max_class_prob;
    if (confidence < conf_threshold)
        continue;

    // bbox
    const float cx = data[i * M + 0];
    const float cy = data[i * M + 1];
    const float cw = data[i * M + 2];
    const float ch = data[i * M + 3];
    // printconst f("cx = %f, cy = %f, w = %f, h = %f\n", cx, cy, w, h);

    // yolov5 输出是相对640的中心坐标和宽高，映射回原图
    const float x = (cx - cw / 2.0f) * scale_x;
    const float y = (cy - ch / 2.0f) * scale_y;
    const float w = cw * scale_x;
    const float h = ch * scale_y;

    pre_boxes.emplace_back(YoloBox {x, y, w, h, confidence, class_id});
  }

  return NMS(pre_boxes, nms_threshold);
}
